import os
import torch
from models.lprload import LPRDataLoader, valuation


def test(model_dir='run_1/best.pth', images_dir=r"../ccpd4/lpr/train", batch_size=128, device=""):
    new_net = True if 'new' in model_dir else False
    test_dataset = LPRDataLoader(os.path.expanduser(images_dir).split(','), new=new_net)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') if not device else device
    lpr_net = torch.load(model_dir, map_location=device)
    print('\n', '网络模型:', 'YOLOv5+LPRNet' if 'ccpd4' in images_dir else '改进的YOLOv5+LPRNet', '测试集:', '蓝色微倾斜' if 'roate' in images_dir else '蓝色强倾斜')
    return valuation(lpr_net, test_dataset, batch_size, device)


if __name__ == '__main__':
    from xlwt import Workbook

    model_files = ['logs/new1/best.pth']
    test_dirs = ['../ccpd8/roate', '../ccpd8/tilt']
    book = Workbook(encoding='utf-8', style_compression=0)
    excel0 = book.add_sheet('ccpd4', cell_overwrite_ok=True)
    excel1 = book.add_sheet('ccpd8', cell_overwrite_ok=True)
    sheet0 = []
    sheet1 = []
    for i in range(1):
        print(f'第{i + 1}次测试')
        for m in model_files:
            for d in test_dirs:
                if 'ccpd4' in d:
                    sheet0.append(test(model_dir=m, images_dir=d))
                else:
                    sheet1.append(test(model_dir=m, images_dir=d))
    for j in range(len(sheet0)):
        for k in range(len(sheet0[0])):
            excel0.write(j, k, sheet0[j][k])
    for j in range(len(sheet1)):
        for k in range(len(sheet1[0])):
            excel1.write(j, k, sheet1[j][k])
    book.save('测试记录.xls')
